[None][feat] Support FP8 KV cache for non-MLA PrimsTS attention - #19396

Description

Enable FP8 E4M3 KV cache for standard non-MLA PrimsTSFmha context, decode, and mixed batches. Fused preprocessing quantizes Q, and PrimTS consumes FP8 Q/K/V while writing native BF16/FP16 output. The effective QK/output scales stay on the GPU and remain live during CUDA Graph replay.

The PrimTS refresh #21354 has merged. This branch is rebased onto main at b2f659aef8437b2fdf9d1cff37a45f5d5eed26e4; the PR Files changed view now contains only the FP8 integration and test changes.

The vendor lock, compatibility patch, and sources remain identical to main: canonical source yuxianq/flashinfer:trtllm-prims-ts-dev-20261107 at 911d6a9d62308b6c8a5249e8dc4aabc7518020f1. Merged FlashInfer #8093 provides device global scales.

Context and decode bind each scale kind with plan(softmax_scale=..., output_scale=...) and pass live values to run(). QK uses preprocessing's natural-unit BMM1 scale at index 0. Upstream infers scalar/tensor kinds independently and includes them in compile identities.

Context preprocessing now clears unused FP8 V rows in the final cache page, using packed FP8 stores in both preprocessing kernels. The existing native-cache behavior is preserved, and PrimTS can use paged_v_tail_is_zero=True without clearing the FP8 tail again.

Standard dense and block-sparse PrimTS share the SM100/SM103/SM107 architecture list. MLA retains its SM100/SM103 routing, and paged block-sparse FP8 KV remains unsupported. Standard dense PrimTS stays opt-in. Performance optimization and default enablement remain follow-up work once performance meets expectations.

The computed-RoPE regression also exposed a one-position shift when the fused preprocessing entry points supplied a zero-length vision region. They now use the disabled-region sentinel (-1); context and decode regressions check the normal RoPE reference.

Validation

Current head fc41ebd016 reduces the five affected GPU test groups from 122 to 27 cases, preserving the nine pre-existing cases. Added component cases fall from 113 to 18. The inexpensive mocked checks, existing MLA/refresh tests, and the single added Qwen3-8B FP8 model-accuracy case remain.

The device-scale check uses one tensor-scale wrapper and an independent PyTorch reference; it no longer compiles two extra float-scale/output-dtype comparison wrappers for every case. Added integration cases compare TRTLLM against the Vanilla reference without separately running the FlashInfer backend.

Current validation: all 27 cases passed on B300, with no failures or skips, using the native build containing this PR's preprocessing fixes and normal library initialization. The retained tail case also fails against the old CI 68751 attention library at the consumer handoff with 15,872 NaN tail elements. That negative control loads the old attention library directly to bypass unrelated obsolete operator registrations. Exact valid-V preservation, zero-tail-before-consumption, independent numerical tolerances, graph replay, and no-host-read assertions are retained. Pre-commit and test collection checks passed.

The following larger GPU and model runs are historical qualification before the matrix reduction and rebase. Full CI passed on head 31e21f57fd in pipeline #68862.

• Built the matching native C++ libraries, Python bindings, and wheel with CUDA 13.4 / PyTorch 2.14.0a0+b2c75dd062.nv26.09 / CuTe DSL 4.8.0. The 114-case attention runs and model evaluation use normal library initialization and verify the freshly built library paths.
• 114 attention cases passed on B300 and the same 114 passed on real Rubin/SM107. This includes 64 poisoned-tail cases across BF16/FP16 input, native/FP8 KV, fresh/cached context, both cache managers, both preprocessing variants, and supported context head dimensions 128/256. The test checks zero tails before the consumer and unchanged valid V values. Four additional cases cover computed-RoPE decode; the broader GQA cases cover context, decode, mixed batches, and long decode.
• 332 additional checks passed on Rubin before the native rebuild (using CI build 68751’s native operations): 80 block-sparse cases, four VisualGen VSA cases (including actual sparse-route assertions and graph replay), four checkpoint-loader regressions, 226 adapter/import checks, four late-key-block cases, 12 FP8 decode scale/graph cases covering every reduction mode and D64/128/256, and two context scale/graph cases.
• Full Qwen3-8B FP8-KV model accuracy passed on B300: 512 CNN/DailyMail samples, score 29.674, unchanged threshold 28.010, reference target 30. Assertions confirmed actual PrimTS context and generation calls.
• Negative control: with the old native library, the poisoned-tail test fails at the consumer handoff with 15,872 NaN tail elements still present. The new tail fixture bounds Q/K to avoid unrelated FP8 rounding-bin noise; its V inputs and exact V checks are unchanged. Existing operator and model-accuracy inputs, references, and tolerances are unchanged.

The local wheel explicitly targets SM103; its common preprocessing kernels target SM100f and ran successfully in the Rubin component tests. Full LLM accuracy was measured on B300. The executor requires an explicitly SM107-enabled build for a full Rubin model run. The Qwen3 non-MLA accuracy gate now allows SM107 for such builds; the MLA gate remains SM100/SM103.

Rebase audit at 17b8fe0c94: no content conflicts; all 12 feature commits were preserved, with the Qwen3.5 checkpoint-loader fix omitted because it is already on main. The vendor lock, patch, and source tree match main. CI helper #83618 was already active on that pre-reduction head when the test-only commit was pushed; a duplicate full run was not requested.

Decode performance

Measured on B300 SXM6 (SM103) at e6df1a889aad, as a baseline before the preprocessing fixes. BF16 activations/output, FP8 E4M3 KV, Hq=28/Hkv=4/D=128, page size 32, KV cache manager v2, fused RoPE, one decode token/request. Both backends use identical warmed inputs. Timings include GPU QKV preprocessing, cache append, attention, and output conversion; they exclude QKV projections, CPU metadata/planning, and JIT compilation.

Method: 16 forwards per CUDA Graph, 20 replays/sample, four ABBA rounds, median of eight samples/backend. Runtime: PyTorch 2.14.0a0+b2c75dd062.nv26.09, CUDA 13.4, FlashInfer 0.7.2, CuTe DSL 4.8.0, real native operations from CI build 68751. No Rubin latency measurement is claimed. These results do not justify enabling standard PrimTS by default in this PR.

PR Checklist

• Description explains FP8 behavior, upstream dependency, and relationship to [None][chore] Refresh PrimTS source onto FlashInfer main (2026-11-07) #21354.
• Tests accompany the feature; numerical references and tolerances are unchanged.
• Vendor source, patch, and lock remain identical to main.
• No new dependency pins or public model/configuration fields.
• Please check this after reviewing the above items as appropriate for this PR.

Normal CI and required reviewer approvals are pending.

GitHub Bot Help

To see available CI bot commands, comment /bot help.

Dev Engineer Review

• PrimsTSFmha adds standard FP8 E4M3 KV-cache support for context and decode on SM100, SM103, and SM107. MLA remains limited to SM100 and SM103. FP8 requests require scalar float32 QK and output scales on the query device; paged block-sparse FP8 remains unsupported.
• Preprocessing supplies FP8 queries and scales to PrimTS. Context and decode wrappers pass tensor scales through to the kernels, including during graph replay. The new query view excludes unused buffer bytes.
• The paged-V tail clear now handles FP8 cache stores as well as matching input and cache types. This change targets NaN propagation from unused rows. Verify packed-store sizing and tail clearing on supported cache layouts.
• The preprocessing API changes vision RoPE defaults from 0 to -1 in three paths. Verify callers that depend on the previous defaults.
• The paged-KV policy helper adds the keyword-only allow_fp8_kv_cache argument. Check external callers if they rely on the previous function signature.
• PrimTS no longer checks the installed CUTLASS DSL package version. It retains compiler-version and operation checks. Verify that supported environments still fail closed when required dependencies are unavailable.
• Standard dense PrimTS remains opt-in. The supplied description reports slower B300 decode timings than trtllm-gen FP8, so this change does not establish a basis for enabling it by default.
• CI, build, and test results for the current head: not available.

QA Engineer Review

The change modifies 14 files: 10 test-code or test-list files and 4 implementation files. The tests add coverage for FP8 support and rejection paths, SM107 routing, V-tail clearing, device-resident scales, graph replay, and preprocessing metadata. The integration entries register BF16- and FP8-KV Qwen3 variants in the B200 test database and manual QA list. Execution results for the current head are not available. Coverage verdict: needs follow-up, pending current-head test results and verification on supported GPUs.

Per-File QA Perspective

• cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h — FP8 cache tails now use packed cache stores to clear unused V rows. Verify valid-row preservation and zeroed tails for FP8 and native cache types.
• cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp — Context preprocessing, context postprocessing, and generation preprocessing now set vision RoPE start and length to -1. Verify behavior for callers that consume these defaults.
• tensorrt_llm/_torch/attention/backends/fmha/cute_dsl_mla.py — MLA page-size exclusivity now uses the MLA-specific supported-SM set. Verify routing remains limited to SM100 and SM103.
• tensorrt_llm/_torch/attention/backends/fmha/prims_ts.py — Adds standard FP8 KV support, SM107 non-MLA routing, device-scale forwarding, and an optional FP8 policy argument. Verify rejection messages, scale validation, query-buffer slicing, and context/decode graph replay.
• tests/unittest/_torch/attention/backend_case.py — Changes prompt-length metadata for context and generation requests. Verify test cases with cached prefixes and mixed context/generation batches. No integration test-list entry applies.
• tests/unittest/_torch/attention/sparse/test_prims_ts_block_sparse.py — Tests SM107 block-sparse selection and FP8 paged-cache rejection. No integration test-list entry applies.
• tests/unittest/_torch/attention/test_cute_dsl_mla.py — Updates the page-size exclusivity expectation for SM107. No integration test-list entry applies.
• tests/unittest/_torch/attention/test_prims_ts_attention_backend.py — Adds coverage for poisoned V tails, computed RoPE, FP8 scaling, graph replay, and decode workspace reuse. No integration test-list entry applies.
• tests/unittest/_torch/attention/test_prims_ts_fmha.py — Adds FP8 support, availability, workspace, preprocessing, and wrapper-scale tests. No integration test-list entry applies.
• tests/unittest/_torch/visual_gen/sparse_attention/vsa/test_vsa_parity.py — Extends parity-test GPU eligibility to SM107. No integration test-list entry applies.
• tests/integration/defs/accuracy/test_llm_api_pytorch.py — Adds BF16- and FP8-KV Qwen3 PrimTS variants and checks the FP8 quantization setting. Both variants are listed in the B200 CI test database and manual QA list.
• tests/integration/defs/accuracy/references/cnn_dailymail.yaml — Adds the Qwen3-8B FP8-KV accuracy reference. Verify it against the relevant evaluation run; the file is a test reference, not a test-list entry.
• tests/integration/test_lists/test-db/l0_b200.yml — Replaces the unparameterized Qwen3 test entry with BF16-KV and FP8-KV variants for CI.
• tests/integration/test_lists/qa/llm_function_core.txt — Replaces the unparameterized Qwen3 entry with BF16-KV and FP8-KV variants for manual QA.